#pragma once
#include "../common/Message.hpp"
#include "../common/Net.hpp"
#include <set>

namespace rpc
{
    namespace server
    {
        class ProviderManager
        {
        public:
            using ptr = std::shared_ptr<ProviderManager>;
            struct ServiceProvider
            {
                using ptr = std::shared_ptr<ServiceProvider>;

                std::mutex _mutex;
                BaseConnection::ptr conn;
                Address host;
                std::vector<std::string> services; // 当前provider可提供的服务

                ServiceProvider(const BaseConnection::ptr &c, const Address &h)
                    : conn(c), host(h)
                {
                }

                void appendService(std::string service) // 追加服务
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    services.push_back(service);
                }
            };

            // 新增服务的提供者
            void addProvider(const BaseConnection::ptr &conn, const Address &host, const std::string &service)
            {
                ServiceProvider::ptr provider;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    // 1. 判断当前连接是否存在
                    auto it = _conn_map_service.find(conn);
                    if (it != _conn_map_service.end())
                    {
                        provider = it->second;
                    }
                    else
                    {
                        provider = std::make_shared<ServiceProvider>(conn, host);
                        _conn_map_service.insert(std::make_pair(conn, provider));
                    }
                    // 2. 建立服务与提供者之间的关联
                    auto &providers = _service_map_provider[service];
                    providers.insert(provider);
                }
                // 3. 服务对象中新增所能提供的服务名称
                provider->appendService(service);
            }

            // 获取连接对应的提供者信息
            ServiceProvider::ptr getProvider(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conn_map_service.find(conn);
                if (it == _conn_map_service.end())
                {
                    return ServiceProvider::ptr();
                }
                return it->second;
            }

            // 删除提供者
            void delProvider(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conn_map_service.find(conn);
                if (it == _conn_map_service.end())
                {
                    return;
                }
                ServiceProvider::ptr provider = it->second;
                for (auto &service : provider->services)
                {
                    auto &providers = _service_map_provider[service];
                    providers.erase(provider);
                }
                _conn_map_service.erase(it);
            }

            std::vector<Address> serviceHosts(const std::string &service)
            {
                auto it = _service_map_provider.find(service);
                if (it == _service_map_provider.end())
                {
                    return std::vector<Address>();
                }
                std::vector<Address> hosts(it->second.begin(), it->second.end());
                return hosts;
            }

        private:
            std::mutex _mutex;
            std::unordered_map<BaseConnection::ptr, ServiceProvider::ptr> _conn_map_service;
            std::unordered_map<std::string, std::set<ServiceProvider::ptr>> _service_map_provider;
        };

        class DiscovererManager
        {
        public:
            using ptr = std::shared_ptr<DiscovererManager>;
            struct ServiceDiscoverer
            {
                using ptr = std::shared_ptr<ServiceDiscoverer>;

                std::mutex _mutex;
                BaseConnection::ptr conn;
                std::vector<std::string> services; // 当前discoverer发现的服务

                ServiceDiscoverer(const BaseConnection::ptr &c) : conn(c)
                {
                }

                void appendService(const std::string &service)
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    services.push_back(service);
                }
            };

            void addDiscoverer(const BaseConnection::ptr &conn, const std::string &service)
            {
                ServiceDiscoverer::ptr discoverer;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it = _conn_map_discoverer.find(conn);
                    if (it != _conn_map_discoverer.end())
                    {
                        discoverer = it->second;
                    }
                    else
                    {
                        discoverer = std::make_shared<ServiceDiscoverer>(conn);
                        _conn_map_discoverer.insert(std::make_pair(conn, discoverer));
                    }
                    auto &discoverers = _service_map_discoverer[service];
                    discoverers.insert(discoverer);
                }
                // 发现者对象新增服务方法名称
                discoverer->appendService(service);
            }

            void delDiscoverer(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conn_map_discoverer.find(conn);
                if (it == _conn_map_discoverer.end())
                {
                    return; 
                }
                auto &discoverer = it->second;
                for (auto &service : discoverer->services)
                {
                    auto &discoverers = _service_map_discoverer[service];
                    discoverers.erase(discoverer);
                }
                _conn_map_discoverer.erase(it);
            }

            void onlineNotify(const std::string &service, const Address& host)
            {
                Notify(service, host, SERVICE_OP_TYPE::SERVICE_ONLINE);
            }

            void offlineNotify(const std::string &service, const Address& host)
            {
                Notify(service, host, SERVICE_OP_TYPE::SERVICE_OFFLINE);
            }
        
        private:
            void Notify(const std::string &service, const Address& host, const SERVICE_OP_TYPE &type)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _service_map_discoverer.find(service);
                if (it->second.empty() == true)
                {
                    return;
                }
                ServiceRequest::ptr msg = MessageFactory::create<ServiceRequest>();
                msg->setId(UUID::uuid());
                msg->setHost(host);
                msg->setMethod(service);
                msg->setMsgType(MSG_TYPE::REQ_SERVICE);
                msg->setOption(type);

                auto &discoverers = it->second;
                for (auto &discover : discoverers)
                {
                    discover->conn->send(msg);
                }
            }
        private:
            std::mutex _mutex;
            std::unordered_map<BaseConnection::ptr, ServiceDiscoverer::ptr> _conn_map_discoverer;
            std::unordered_map<std::string, std::set<ServiceDiscoverer::ptr>> _service_map_discoverer;
        };

        // Provider_Discoverer_Manager
        class PDManager
        {
        public:
            using ptr = std::shared_ptr<PDManager>;
            PDManager()
                : _providers(std::make_shared<ProviderManager>())
                , _discoverers(std::make_shared<DiscovererManager>())
            {}

            void onServiceRequest(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                // 服务注册 or 服务发现
                std::unique_lock<std::mutex> lock(_mutex);
                auto type = msg->Option();
                if (type == SERVICE_OP_TYPE::SERVICE_REGISTER)
                {
                    // 1. 服务注册
                    _providers->addProvider(conn, msg->host(), msg->Method()); 
                    // 2. 服务上线通知
                    _discoverers->onlineNotify(msg->Method(), msg->host());
                    registryResponse(conn, msg);
                }
                else if (type == SERVICE_OP_TYPE::SERVICE_DISCOVERY)
                {
                    _discoverers->addDiscoverer(conn, msg->Method());
                    discoveryResponse(conn, msg);
                }
                else
                {
                    ELOG("receive service request, service type error");
                    errorResponse(conn, msg);
                }
            }

            void onConnShutDown(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _providers->getProvider(conn);
                if (it.get() != nullptr) // conn是服务提供者的连接
                {
                    // 服务下线通知
                    for (auto &service : it->services)
                    {
                        _discoverers->offlineNotify(service, it->host);
                    }
                    _providers->delProvider(conn);
                    return;
                }
                _discoverers->delDiscoverer(conn);
            }
        private:
            // 针对服务提供者的服务注册响应
            void registryResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &req)
            {
                auto msg_rsp = MessageFactory::create<ServiceResponse>();
                msg_rsp->setId(req->id());
                msg_rsp->setMsgType(MSG_TYPE::RSP_SERVICE);
                msg_rsp->setRcode(RCODE::RCODE_OK);
                msg_rsp->setOption(SERVICE_OP_TYPE::SERVICE_REGISTER);
                
                conn->send(msg_rsp);
            }
            // 针对服务发现者的服务发现响应
            void discoveryResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &req)
            {
                auto msg_rsp = MessageFactory::create<ServiceResponse>();
                msg_rsp->setId(req->id());
                msg_rsp->setOption(SERVICE_OP_TYPE::SERVICE_DISCOVERY);
                msg_rsp->setMsgType(MSG_TYPE::RSP_SERVICE);

                std::string method = req->Method();
                auto hosts = _providers->serviceHosts(method);
                if (hosts.empty() == true)
                {
                    msg_rsp->setRcode(RCODE::RCODE_NOT_FOUND_SERVICE);
                    conn->send(msg_rsp);
                    return;
                }
                msg_rsp->setRcode(RCODE::RCODE_OK);
                msg_rsp->setMethod(req->Method());
                msg_rsp->setHosts(hosts);
                conn->send(msg_rsp);
            }
            
            // 错误响应
            void errorResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &req)
            {
                auto msg_rsp = MessageFactory::create<ServiceResponse>();
                msg_rsp->setId(req->id());
                msg_rsp->setMsgType(MSG_TYPE::RSP_SERVICE);
                msg_rsp->setOption(SERVICE_OP_TYPE::SERVICE_UNKNOWN);
                msg_rsp->setRcode(RCODE::RCODE_ERROR_MSGTYPE);
                conn->send(msg_rsp);
            }
        private:
            std::mutex _mutex;
            ProviderManager::ptr _providers;
            DiscovererManager::ptr _discoverers;
        };
    }
}